import dspy
import logging
from openai import OpenAI

# 豆包模型API配置
DOUBAO_API_BASE = "https://ark.cn-beijing.volces.com/api/v3"
DOUBAO_MODEL = "doubao-seed-1-6-flash-250615"

# 配置日志记录
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(),
        logging.FileHandler("dspy_doubao.log")
    ]
)
logger = logging.getLogger("DoubaoModel")


class DoubaoLM(dspy.LM):
    def __init__(self, api_key: str, model: str = DOUBAO_MODEL):
        super().__init__(model)
        self.model = model
        self.client = OpenAI(
            base_url=DOUBAO_API_BASE,
            api_key=api_key,
        )

    def forward(self, prompt=None, messages=None, **kwargs):
        # 确保至少有一个消息
        if not messages:
            raise ValueError("No messages provided for the chat completion.")

        print("="*50)
        for msg in messages:
            print(f"\n====> {msg['role']}: {msg['content']}")

        # 调用豆包API
        return self.client.chat.completions.create(
            model=self.model,
            messages=messages,
            **kwargs
        )
